package com.peng.leetcode.heap;

import java.util.Iterator;
import java.util.Map;
import java.util.TreeMap;

/**
 * GetLeastNumbers
 *
 * @author: lupeng6
 * @create: 2020/12/11 17:28
 */
public class GetLeastNumbers {

    public static void main(String[] args) {
        int[] arr = {0, 0, 1, 2, 4, 2, 2, 3, 1, 4};
        arr = new Solution().getLeastNumbers(arr, 8);
        for (int i : arr) {
            System.out.println(i);
        }
    }

    static class Solution {
        public int[] getLeastNumbers(int[] arr, int k) {
            TreeMap<Integer, Integer> treeMap = new TreeMap<>();
            for (int i : arr) {
                Integer integer = treeMap.get(i);
                if (integer != null) {
                    integer++;
                    treeMap.put(i, integer);
                } else {
                    treeMap.put(i, 1);
                }
            }
            int[] retArr = new int[k];

            Iterator<Map.Entry<Integer, Integer>> iterator = treeMap.entrySet().iterator();
            int i = 0;
            while (iterator.hasNext() && i < k) {
                Map.Entry<Integer, Integer> entry = iterator.next();
                Integer value = entry.getValue();
                for (Integer integer = 0; integer < value && i < k; integer++) {
                    retArr[i] = entry.getKey();
                    i++;
                }
            }
            return retArr;
        }
    }
}
